package 困难.并查集;

import java.util.Arrays;

/**
 * 给你一个 n 个点组成的无向图边集 edgeList ，其中 edgeList[i] = [ui, vi, disi] 
 * 表示点 ui 和点 vi 之间有一条长度为 disi 的边。请注意，两个点之间可能有 超过一条边 。
 * 给你一个查询数组queries ，其中 queries[j] = [pj, qj, limitj] ，你的任务是对
 * 于每个查询 queries[j] ，判断是否存在从 pj 到 qj 的路径，且这条路径上的每一条边都 严格小于 limitj 。
 * 请你返回一个 布尔数组 answer ，其中 answer.length == queries.length ，当 queries[j] 
 * 的查询结果为 true 时， answer 第 j 个值为 true ，否则为 false 。
 * <p>
 * 来源：力扣（LeetCode）
 * 链接：https://leetcode.cn/problems/checking-existence-of-edge-length-limited-paths
 */
public class 检查边长度限制的路径是否存在_1697 {

    public static void main(String[] args) {

    }

    /**
     * 并查集+离线查询
     * 1，将无向图按照边长度从小到大排序
     * 2，将查询按照 limit 排序
     * 3，使用并查集将小于当前查询 limit 的所有 edge 全部连通起来
     */
    public boolean[] distanceLimitedPathsExist(int n, int[][] edgeList, int[][] queries) {
        boolean[] ans = new boolean[queries.length];

        int[][] q = new int[queries.length][4];
        for (int i = 0; i < q.length; i++) {
            q[i] = new int[]{i, queries[i][0], queries[i][1], queries[i][2]};
        }
        Arrays.sort(q, (a, b) -> a[3] - b[3]);
        Arrays.sort(edgeList, (a, b) -> a[2] - b[2]);

        DSU dsu = new DSU(n + 1);

        int edgeIdx = 0;
        for (int i = 0; i < q.length; i++) {

            for (; edgeIdx < edgeList.length && edgeList[edgeIdx][2] < q[i][3]; edgeIdx++) {
                dsu.union(edgeList[edgeIdx][0], edgeList[edgeIdx][1]);
            }
            ans[q[i][0]] = dsu.connected(q[i][1], q[i][2]);
        }
        return ans;
    }

    static class DSU {

        private int[] pre;

        public DSU(int n) {
            pre = new int[n];
            for (int i = 0; i < pre.length; i++) {
                pre[i] = i;
            }
        }

        public int find(int x) {
            if (pre[x] == x) {
                return x;
            }
            return pre[x] = find(pre[x]);
        }

        public void union(int x, int y) {
            x = find(x);
            y = find(y);
            if (x == y) {
                return;
            }
            pre[x] = y;
        }

        public boolean connected(int x, int y) {
            return find(x) == find(y);
        }
    }

}
